import argparse
import datetime
import os
import sys

import cv2
import numpy as np

parse = argparse.ArgumentParser(prog="take_photo", prefix_chars="-", description="Press [space] to save an image", epilog="Miaow")

parse.add_argument('-o', '--output', default="./output", help="The folder photos saved")
parse.add_argument("-i", "--info", action="store_true", help="Show camera info")

args = parse.parse_args()
os.makedirs(args.output, exist_ok=True)
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
time_stamp = datetime.datetime.now().strftime("%Y%m%dT%H%M%S")
cap = cv2.VideoCapture(0)
if not cap.isOpened():
    print("Not Opened")
    sys.exit(1)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 480)
if args.info:
    print("========Camera info==========")
    print(cap.get(cv2.CAP_PROP_FRAME_WIDTH), "x", cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    print(cap.get(cv2.cv2.CAP_PROP_FPS), "fps")
    print(cap.get(cv2.cv2.CAP_PROP_EXPOSURE), "exposure")
    print("=============================")

cv2.namedWindow("realtime display")
img_num = 0
count_down = 4
key_code = 0
text = "Press [SPACE]"
while True:
    status, img = cap.read()
    img_raw = img.copy()
    if not status:
        continue
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # Find the chess board corners
    ret, corners = cv2.findChessboardCorners(gray, (9, 6), None, cv2.CALIB_CB_FAST_CHECK)
    if ret == True:
        corners2 = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
        img_chessboard = cv2.drawChessboardCorners(img, (9, 6), corners2, ret)
    else:
        img_chessboard = img

    if count_down < 4:
        count_down = count_down + 1
        img_display = np.uint8(np.clip((img_chessboard / 3 * (4 - count_down) + 45 * count_down), 0, 255))  # 1.5是对比度   10是亮度

    else:
        img_display = img_chessboard

    img_display = cv2.putText(cv2.flip(img_display, 0), text, (10, 25), fontFace=cv2.FONT_HERSHEY_PLAIN, fontScale=1, color=(255, 255, 255), thickness=1)
    if key_code == ord("q"):
        break
    elif key_code == ord(" "):
        img_num = img_num + 1
        count_down = 0
        cv2.imwrite(os.path.join(args.output, str(img_num) + "_" + time_stamp + ".png"), img_raw)
        text = str(img_num) + "_" + time_stamp + ".png saved"

    cv2.imshow("realtime display", img_display)
    key_code = cv2.waitKey(5)

cap.release()
cv2.destroyAllWindows()
